package ekforth

import java.io.Closeable
import java.io.IOException
import java.net.ServerSocket
import kotlin.concurrent.thread

class StackUnderflowException : Exception()

class MemoryCell(
    val getter: () -> Any?,
    val setter: (Any?) -> Unit
)

class Stack {
    val stack: MutableList<Any?> = mutableListOf()

    fun push(obj: Any?) {
        stack.add(obj)
    }

    fun pop(): Any? = if (stack.isEmpty()) {
        throw StackUnderflowException()
    } else {
        stack.removeAt(stack.size - 1)
    }

    fun peek() = if (stack.isEmpty()) {
        throw StackUnderflowException()
    } else {
        stack.last()
    }

    fun clear() {
        stack.clear()
    }
}

class Dictionary {
    val entries: MutableList<Any?> = mutableListOf()
    operator fun get(i: Int) = entries[i]
    operator fun set(i: Int, value: Any?) {
        entries[i] = value
    }
    fun add(obj: Any?): Int {
        entries.add(obj)
        return here() - 1
    }
    fun here(): Int = entries.size

    override fun toString() = entries.withIndex().map { "${it.index}: ${it.value}" }.joinToString("\n") { it }

}

private sealed class Instruction {
    abstract fun perform(forth: EKForth): String?
    class FunCall(val function: (EKForth) -> Unit) : Instruction() {
        override fun perform(forth: EKForth): String? {
            ++ forth.eip
            try {
                function.invoke(forth)
                return null
            } catch(e: Exception) {
                return e.toString()
            }
        }

        override fun toString(): String = "FUNCALL"
    }

    class Jump(val address: Int) : Instruction() {
        override fun perform(forth: EKForth): String? {
            if (address < forth.dictionary.here()) {
                forth.eip = address
                return null
            } else {
                return("Invalid address")
            }
        }
        override fun toString(): String = "JUMP $address"
    }
    class DynJump(val address: () -> Int) : Instruction() {
        override fun perform(forth: EKForth): String? {
            val addr = address.invoke()
            if (addr < forth.dictionary.here()) {
                forth.eip = addr
                return null
            } else {
                return("Invalid address")
            }
        }

        override fun toString(): String = "JUMP $address"
    }
    object End : Instruction() {
        override fun perform(forth: EKForth): String? {
            forth.end = true
            return null
        }
        override fun toString(): String = "END"
    }

}

sealed class Input {

    companion object {
        private val eof = ForthWord("")
        fun accept(string: String) = StringInput(string)
        fun accept(input: Iterator<Any?>) = ParsedInput(input)
        fun accept(input: List<Any?>) = ParsedInput(input.iterator())
    }

    abstract fun word(): Any?

    class StringInput(val string: String) : Input() {
        var start = 0
        override fun word(): Any? {
            while (start < string.length && string[start].isWhitespace()) ++start
            if (start == string.length) return eof
            val tokenStart = start
            while (start < string.length && !string[start].isWhitespace()) ++start
            val token = string.substring(tokenStart, start)
            return token.toIntOrNull() ?: token.toDoubleOrNull() ?: ForthWord(token)
        }
    }

    class ParsedInput(val input: Iterator<Any?>) : Input() {
        override fun word(): Any? = if (input.hasNext()) input.next() else eof
    }
}

@Suppress("UNCHECKED_CAST")
class EKForth {

    var lastDefinition: Int = -1
    val dictionary = Dictionary()
    val stack = Stack()
    val returnStack = Stack()
    var ip = 0
    var w: Int = 0
    var eip = -1
    var end = true
    var pfa = -1

    data class Pointers(
        val eip: Int,
        val ip: Int,
        val w: Int,
        val pfa: Int
    )

    var input: Input = Input.accept("")

    fun accept(input: String) {
        this.input = Input.accept(input)
        run()
    }
    fun accept(input: List<Any?>) {
        this.input = Input.accept(input)
        run()
    }
    fun accept(input: Iterator<Any?>) {
        this.input = Input.accept(input)
        run()
    }


    val next = dictionary.here().also {
        dictionary.add(Instruction.FunCall { forth ->
            with (forth) {
                w = dictionary[ip] as Int;
                ++ ip
                pfa = w + 1
            }
        } )
        dictionary.add(Instruction.DynJump { w })
    }

    private val docreate = dictionary.here().also {
        dictionary.add(Instruction.FunCall { forth ->
            with(forth) {
                stack.push(pfa)
            }
        })
        dictionary.add(Instruction.Jump(next))
    }

    private val dovar = docreate

    val state = createHeader("STATE").also {
        comma(Instruction.Jump(dovar))
        comma(0)
    }.also { lastDefinition = it }

    val variable = defWord("VARIABLE") { forth ->
        with (forth) {
            lastDefinition = createHeader(((kWord() as ForthWord).word))
            comma(Instruction.Jump(dovar))
            comma(0)
        }
    }

    val enter = dictionary.here().also {
        dictionary.add(Instruction.FunCall { forth ->
            with(forth) {
                returnStack.push(ip)
                ip = w + 1
            }
        })
        dictionary.add(Instruction.Jump(next))
    }

    val exit = dictionary.here().also {
        dictionary.add(Instruction.FunCall { forth -> forth.ip = forth.returnStack.pop() as Int })
        dictionary.add(Instruction.Jump(next))
    }

    val eoln = defWord("") { forth -> forth.ip = forth.returnStack.pop() as Int }


    /*
    header:
    - link
    - precedence
    - name
     */
    private fun defWord(name: String, action: Instruction, precedence: Boolean): Int {
        createHeader(name, precedence)
        lastDefinition = dictionary.here()
        dictionary.add(action)
        return lastDefinition
    }

    fun createHeader(name: String, precedence: Boolean = false): Int {
        dictionary.add(lastDefinition)
        dictionary.add(precedence)
        dictionary.add(name)
        return dictionary.here()
    }

    fun definitionLink(address: Int): Int = address - 3
    fun definitionName(address: Int): Int = address - 1
    fun definitionPrecedence(address: Int): Int = address - 2
    fun definitionCodeField(address: Int): Int = address

    fun defWord(name: String, action: (EKForth) -> Unit) =
        defWord(
            name = name,
            action = Instruction.FunCall(action),
            precedence = false
        ).also {
            dictionary.add(Instruction.Jump(next))
        }

    fun findWord(name: String): Int? {
        var def = lastDefinition
        while (def >= 0) {
            if (name.equals(dictionary[definitionName(def)] as String, ignoreCase = true)) return definitionCodeField(def)
            def = dictionary[definitionLink(def)] as Int
        }
        return null
    }

    var output: (Any?) -> Unit = ::print

    val dot = defWord(".") { forth ->
        output.invoke(forth.stack.pop())
    }

    val docon = defWord("DOCON") { forth ->
        forth.stack.push(forth.dictionary[forth.pfa])
    }

    // ;C EXECUTE   i*x xt -- j*x   execute Forth word
    val execute = createHeader("EXECUTE").also {
        dictionary.add(Instruction.FunCall { forth ->
            with (forth) {
                w = stack.pop() as Int;
                pfa = w + 1
                //++ ip
            }
        }).also { lastDefinition = it }
        dictionary.add(Instruction.DynJump { w })
    }


    val kaccept = defWord("KACCEPT") { forth ->
        with (forth) {
            save()
            end = true
        }
    }

    fun defColon(name: String, vararg tokens: Any) = defWord(name, Instruction.Jump(enter), precedence = false).also {
        tokens.forEach { xt -> dictionary.add(xt) }
        dictionary.add(exit)
    }

    fun colonDefinition(name: String, precedence: Boolean = false, definition: () -> Unit) = createHeader(name, precedence).also {
        comma(Instruction.Jump(enter))
        definition.invoke()
        comma(exit)
        lastDefinition = it
    }

    val bye = defWord("BYE", Instruction.End, false)

    val onStack = defWord("ONSTACK") { forth ->
        with (forth) {
            stack.push(dictionary[ip++])
        }
    }

    val branch = defWord("BRANCH") { forth ->
        forth.ip = dictionary[forth.ip] as Int
    }

    val qbranch = defWord("?BRANCH") { forth ->
        if (stack.pop() == 0) {
            forth.ip = dictionary[forth.ip] as Int
        } else {
            ++ forth.ip
        }
    }

    fun kWord() = input.word()

    val kWord = defWord("KWORD") { forth ->
        forth.stack.push(kWord())
    }

    val kFind = defWord("KFIND") { forth ->
        with (forth) {
            val token = stack.pop()
            if (token is ForthWord) {
                val definition = findWord(token.word) ?: throw Exception("$token not found")
                val isImmediate = if (dictionary[definitionPrecedence(definition)] as Boolean) 1 else -1
                val xt = definitionCodeField(definition)
                stack.push(xt)
                stack.push(isImmediate)
            } else {
                stack.push(token)
                stack.push(0)
            }
        }
    }

    private fun begin() {
        stack.push(dictionary.here())
    }

    private fun again() {
        comma(branch)
        comma(stack.pop())
    }

    val dup = defWord("DUP") { forth -> forth.stack.push(forth.stack.peek()) }
    val qdup = defWord("?DUP") { forth ->
        val n = forth.stack.peek()
        if (n != 0) forth.stack.push(n)
    }
    val drop = defWord("DROP") { forth -> forth.stack.pop() }

    val swap = defWord("SWAP") { forth ->
        val x = stack.pop()
        val y = stack.pop()
        stack.push(x)
        stack.push(y)
    }

    val over = defWord("OVER") { forth ->
        val x = stack.pop()
        val y = stack.peek()
        stack.push(x)
        stack.push(y)
    }

    val nip = defWord("NIP") { forth ->
        val x = stack.pop()
        stack.pop()
        stack.push(x)
    }

    val twoDup = defWord("2DUP") { forth ->
        val x = stack.pop()
        val y = stack.peek()
        stack.push(x)
        stack.push(y)
        stack.push(x)
    }

    val comma = defWord(",") { forth -> forth.dictionary.add(forth.stack.pop()) }
    val compilecomma = defWord("COMPILE,") { forth -> forth.dictionary.add(forth.stack.pop()) }
    val compileLiteral = defWord("COMPILE-LITERAL") { forth ->
        with (forth) {
            val literal = stack.pop()
            if (literal is ForthWord) error("Cannot inline ForthWord")
            forth.dictionary.add(onStack)
            forth.dictionary.add(literal)
        }
    }

    private fun comma(obj: Any?): Int = dictionary.add(obj)

    val zerolessthan = defWord("0<") { forth ->
        val n = forth.stack.pop()
        val flag = if (n is Int) { n > 0 } else { n as Double > 0 }
        forth.stack.push(if (flag) -1 else 0)
    }

    val zerogreaterthan = defWord("0>") { forth ->
        val n = forth.stack.pop()
        val flag = if (n is Int) { n < 0 } else { (n as Double) < 0 }
        forth.stack.push(if (flag) -1 else 0)
    }

    val zeroequal = defWord("0=") { forth ->
        forth.stack.push(if (forth.stack.pop() == 0) -1 else 0)
    }


    val rBracket = defWord("]") { forth ->
        dictionary[state + 1] = -1
    }

    val semicolon = defWord(";") { forth ->
        with(forth) {
            comma(exit)
            lastDefinition = stack.pop() as Int
            dictionary[state + 1] = 0
        }
    }.also {
        dictionary[definitionPrecedence(it)] = true
    }

    val immediate = defWord("IMMEDIATE") { forth ->
        dictionary[definitionPrecedence(lastDefinition)] = true
    }

    val colon = defWord(":") { forth ->
        with (forth) {
            stack.push(createHeader((kWord() as ForthWord).word))
            dictionary.add(Instruction.Jump(enter))
            dictionary[state + 1] = -1
        }
    }

    private fun fif() {
        comma(qbranch)
        stack.push(dictionary.here())
        comma(0)
    }

    private fun felse() {
        dictionary[stack.pop() as Int] = dictionary.here() + 2
        comma(branch)
        stack.push(dictionary.here())
        comma(0)
    }

    fun fthen() {
        dictionary[stack.pop() as Int] = dictionary.here()
    }

    val createHeader = defWord("CREATEHEADER") { forth ->
        forth.defWord((forth.stack.pop() as ForthWord).word, action = Instruction.Jump(forth.enter), precedence = false)
    }

    val create = defColon("CREATE", kWord, createHeader)

    val clearReturnStack = defWord("CLEAR-RETURN-STACK") { forth ->
        forth.returnStack.clear()
    }

    val printOK = defWord(".OK") {
        output.invoke("   ok ")
    }

    val cr = defWord("CR") {
        output.invoke("\n")
    }

    val fetch = defWord("@") { forth ->
        with (forth) {
            val value = when (val cell = stack.pop()) {
                is MemoryCell -> cell.getter.invoke()
                is Int -> dictionary[cell]
                else -> throw Exception("Invalid cell $cell")
            }
            stack.push(value)
        }
    }

    val store = defWord("!") { forth ->
        with (forth) {
            val cell = stack.pop()
            val value = stack.pop()
            when (cell) {
                is MemoryCell -> cell.setter.invoke(value)
                is Int -> dictionary[cell] = value
                else -> throw Exception("Invalid cell $cell")
            }
        }
    }

    val begin = defWord("BEGIN") { forth -> forth.stack.push(forth.dictionary.here().also {println("begin $it")}) }.also {
        dictionary[definitionPrecedence(it)] = true
    }

    val again = defWord("AGAIN") { forth ->
        with (forth) {
            dictionary.add(branch)
            dictionary.add(stack.pop())
        }
    }.also {
        dictionary[definitionPrecedence(it)] = true
    }

    val interpret = colonDefinition("INTERPRET") {
        begin()
        comma(kWord)
        comma(kFind)
        comma(qdup)
        fif()
        comma(state)
        comma(fetch)
        fif()
        comma(zerolessthan)
        fif()
        comma(execute)
        felse()
        comma(compilecomma)
        fthen()
        felse()
        comma(drop)
        comma(execute)
        fthen()
        felse()
        comma(state)
        comma(fetch)
        fif()
        comma(compileLiteral)
        fthen()
        //comma(wordNotFoundException)
        fthen()
        again()
    }

    val setInterpretationState = defWord("SET-INTERPRETATION-STATE") { forth ->
        with (forth) {
            dictionary[state + 1] = 0
        }
    }

    val quit = colonDefinition("QUIT") {
        begin()
        comma(clearReturnStack)
        comma(setInterpretationState)
        comma(kaccept)
        comma(interpret)
        comma(printOK)
        comma(cr)
        again()
        comma(bye)
    }

    val start = dictionary.here().also {
        dictionary.add(quit)
        dictionary.add(bye)
    }

    private val initPointers = Pointers(next, start, 0, 0)
    private var pointers = initPointers

    init {
        run()
    }

    fun init() {
        eip = pointers.eip
        ip = pointers.ip
        w = pointers.w
        pfa = pointers.pfa
    }

    fun save() {
        pointers = Pointers(eip, ip, w, pfa)
    }

    fun mainLoop() {
        while (!end) {
            //println("eip: $eip ip: $ip w: $w pfa: $pfa stack: ${stack.stack} return stack: ${returnStack.stack}")
            val error = (dictionary[eip] as Instruction).perform(this)
            if (error != null) {
                output.invoke("$error\n")
                pointers = initPointers
                run()
            }
        }
    }

    fun run() {
        init()
        end = false
        mainLoop()
    }

    val dotS = defWord(".S") { forth ->
        output.invoke(stack.stack.map { it.toString() }.joinToString(" "))
    }

    val mul = defWord("*") { forth ->
        val x = forth.stack.pop()
        val y = forth.stack.pop()
        val result = if (x is Int) {
            if (y is Int) {
                x * y
            } else {
                x * (y as Double)
            }
        } else {
            (x as Double) * (y as Double)
        }
        forth.stack.push(result)
    }

    val add = defWord("+") { forth ->
        val x = forth.stack.pop()
        val y = forth.stack.pop()
        val result = if (x is Int) {
            if (y is Int) {
                x + y
            } else {
                x + (y as Double)
            }
        } else {
            (x as Double) + (y as Double)
        }
        forth.stack.push(result)
    }

    val sub = defWord("-") { forth ->
        val y = forth.stack.pop()
        val x = forth.stack.pop()
        val result = if (x is Int) {
            if (y is Int) {
                x - y
            } else {
                x - (y as Double)
            }
        } else {
            (x as Double) - (y as Double)
        }
        forth.stack.push(result)
    }

    val div = defWord("/") { forth ->
        val y = forth.stack.pop()
        val x = forth.stack.pop()
        val result = if (x is Int) {
            if (y is Int) {
                x / y
            } else {
                x / (y as Double)
            }
        } else {
            (x as Double) / (y as Double)
        }
        forth.stack.push(result)
    }

    val negate = defWord("NEGATE") { forth ->
        val x = forth.stack.pop()
        val result = if (x is Int) {
            - x
        } else {
            - (x as Double)
        }
        forth.stack.push(result)
    }

    val q = colonDefinition("?") {
        comma(fetch)
        comma(dot)
    }

}

data class ForthWord(val word: String)

val mainFn = ::forth
//val mainFn = { remoteForth(12345) }

fun main() {
    mainFn.invoke()
}

fun forth() {
    EKForth().run {
        output = { print(it) }
        while (true) {
            val line = readLine() ?: break
            accept(line)
        }
    }
}

fun remoteForth(port: Int) = ForthIOServer(EKForth(), port)

class ForthIOServer(val forth: EKForth, port: Int, backlog: Int = 5) : Closeable {

    private val oldOutput = forth.output
    private val sock = ServerSocket(port, backlog)
    private var thread: Thread = thread {
        while (!Thread.currentThread().isInterrupted) {
            val client = sock.accept()
            val reader = client.getInputStream().bufferedReader()
            val writer = client.getOutputStream().bufferedWriter()
            forth.output = { writer.write("$it"); writer.flush() }
            while (true) {
                val input = reader.readLine() ?: break
                try {
                    forth.accept(input)
                } catch (e: IOException) {
                    break
                }
            }
            forth.output = oldOutput
        }
    }

    override fun close() {
        forth.output = oldOutput
        sock.close()
        thread.interrupt()
    }
}
